# utf-8
import click
from PyPDF2 import PdfReader, PdfWriter, PdfMerger
import glob
import os
import shutil
import time
from tqdm import tqdm


def prepare_output_dir(output_path):
    """确保输出目录存在（若存在则先删除）"""
    if os.path.exists(output_path):
        # 递归删除目录
        shutil.rmtree(output_path)

        # 创建目录
    os.makedirs(output_path, exist_ok=True)


def validate_pdf_extension(ctx, param, value):
    """验证输出路径必须以 .pdf 结尾"""
    if not value.lower().endswith('.pdf'):
        raise click.BadParameter("输出文件必须是 .pdf 格式")
    return value


@click.group()
def cli():
    """PDF 工具箱(吾图南)"""
    pass


@cli.command()
@click.option("--input", "-i", required=True, type=click.Path(exists=True), help="输入PDF文件")
@click.option("--output", "-o", default="./out", type=click.Path(), help="输出目录 (默认: ./out)")
@click.option("--per", "-p", default=5, type=int, help="每份页数 (默认: 5)")
def split(input, output, per):
    """split 子命令 分割PDF文件"""
    """
    将PDF文件按指定页数分割成多个小文件

    参数：
        input: 输入PDF文件路径
        output: 输出文件夹路径
        per: 每个分割文件包含的页数（默认为5）
    """
    # 确保输出文件夹存在
    prepare_output_dir(output)
    # 获取输入文件的名字
    # python main.py split -i ./1.pdf
    # python main.py split -i 1.pdf
    file_name = os.path.basename(input)
    # 读取原始PDF文件
    reader = PdfReader(input)
    total_pages = len(reader.pages)
    # 计算需要分割成多少份
    num_splits = (total_pages + per - 1) // per
    click.secho(f"[{file_name}] 总页数: {total_pages}", fg='blue')
    click.secho(f"[{file_name}] 将分割为 {num_splits} 个文件，每个文件 {per} 页", fg='blue')
    total = num_splits
    # 第一层for 分割PDF
    for i in range(num_splits):

        # 计算进度条长度和百分比
        progress = i // 2
        remaining = (total - i) // 2
        # 多少个循环
        percent = i

        # 使用 \r 回到行首，动态更新同一行
        click.secho(
            f"\r进度: {i + 1}/{total} [{'#' * progress}{' ' * remaining}] {percent + 1}/{total}",
            fg="green",
            nl=False
        )
        # 输出的第一步： pdf 输出器对象
        writer = PdfWriter()
        # 计算当前分割的起始和结束页
        start_page = i * per
        end_page = min((i + 1) * per, total_pages)
        # 添加页到当前分割文件
        for page_num in range(start_page, end_page):
            # 添加当前页到要输出的集合里面
            writer.add_page(reader.pages[page_num])
        # 生成输出文件名。用户输出的路径 与 动态路径 拼接生成新路径。
        output_filename = os.path.join(output, f"{i + 1}.pdf")
        # ./out\65.pdf 替换输出为 .\out\65.pdf(windows的风格)
        replace_output = output_filename.replace("/", "\\")
        # 写入输出文件
        with open(replace_output, "wb") as out:
            writer.write(out)
        click.secho(f" > 已创建: {replace_output} (页 {start_page + 1}-{end_page})", fg="yellow")
    # len(reader.pages) 总页数 // 切割的份数 + 1
    click.secho(f"✓ 已分割为 {len(reader.pages) // per + 1} 个文件", fg="red")


@cli.command()
@click.option("--files", "-f", multiple=True, type=click.Path(exists=True))
@click.option("--input-dir", "-i", type=click.Path(exists=True))
@click.option("--output", "-o", required=True, type=click.Path(writable=True), callback=validate_pdf_extension,
              help="输出路径（必须以 .pdf 结尾）")
def join(files, input_dir, output):
    """合并PDF文件（示例功能，实际需调整）"""
    # files 是元组
    # print(files)
    # print(bool(files))
    # print(type(files))
    # -f 与 -i 二选一
    if not files and not input_dir:
        raise click.UsageError("必须指定 -f 或 -i")

    pdf_files = []
    estimated_size = 0
    if input_dir:
        pdf_files = glob.glob(os.path.join(input_dir, "*.pdf"))
        # 返回的是list
        # print(type(pdf_files))
        # files = list(files) + glob.glob(f"{input_dir}/*.pdf")
        estimated_size: int = sum(os.path.getsize(file) for file in pdf_files)
        # 需要对list进行排序，不然合并文件顺序不对
        pdf_files = sorted(pdf_files, key=lambda x: int(x.split('\\')[-1].split('.')[0]))
        click.secho(f"{input_dir}：目录总大小：{estimated_size} Bytes", fg='blue')

    # 用户输入单个文件
    if files:
        pdf_files = list(files)
        estimated_size = sum(os.path.getsize(file) for file in pdf_files)
        click.secho(f"用户输入文件的总大小：{estimated_size} Bytes", fg='blue')

    # 预估总大小

    # for f in pdf_files:
    #     print(f'name:{f} size: {os.path.getsize(f)}')

    # 2. 合并PDF
    merger = PdfMerger()
    with click.progressbar(pdf_files,
                           label=click.style(f"合并：{len(pdf_files)}个文件", fg="yellow"),
                           bar_template='%(label)s  [%(bar)s]  %(info)s',
                           fill_char=click.style('#', fg='yellow'),
                           ) as bar:
        for pdf in bar:
            merger.append(pdf)

    with tqdm(total=100, desc="输出合并的PDF...") as pbar:
        # 自动缓慢更新到 90%（模拟写入前的准备）
        for _ in range(90):
            time.sleep(0.02)  # 调整 sleep 时间控制速度
            pbar.update(1)  # 每次 +1%

        # 3. 输出到指定路径
        merger.write(output)

        # 4. 手动直接跳到 100%（表示完成）
        pbar.n = 100  # 直接设置当前值为 100
        pbar.refresh()  # 强制刷新显示

        merger.close()
        click.secho(f"\n✓ 合并完成: {len(pdf_files)}个文件 -> {output}", fg="green")


if __name__ == "__main__":
    cli()
